import pandas as pd
import warnings
from TradeSignalCalculator import TradeSignalCalculator
from sqlalchemy import create_engine


warnings.filterwarnings('ignore')

class TradeProcessor:
    def __init__(self, db_connection_string=None, df=None):
        """
        初始化策略因子处理器
        :param db_connection_string: MySQL数据库连接字符串，例如 'mysql+pymysql://user:password@host:port/dbname'
        :param df: 直接传入的DataFrame
        """
        if db_connection_string:
            # 创建数据库连接
            engine = create_engine(db_connection_string)
            query = "SELECT * FROM stock_basic_indicators"
            self.data = pd.read_sql(query, engine)
            engine.dispose()
        elif df is not None:
            self.data = df.copy()
        else:
            raise ValueError("必须提供db_connection_string或df参数")
        self._preprocess_data()
    
    def _preprocess_data(self):
        """数据预处理"""
        # 首先确保 'date' 字段可以被正确解析为日期
        try:
            self.data['date'] = pd.to_datetime(self.data['date'], errors='coerce')  # 如果无法解析，将会变为 NaT
        except Exception as e:
            raise ValueError(f"日期转换错误: {e}")
        
        # 检查是否有 NaT 值，并去除无效日期
        if self.data['date'].isna().sum() > 0:
            print(f"警告：存在无法解析的日期，已被转换为 NaT，数量：{self.data['date'].isna().sum()}")
            self.data = self.data.dropna(subset=['date'])  # 删除 NaT 值的行

        # 进一步处理日期
        self.data['date'] = self.data['date'].dt.normalize().dt.tz_localize(None)  # 规范化日期，并去除时区信息
        
        # 增加 datekey 字段，格式为 'yyyyMMdd'
        self.data['datekey'] = self.data['date'].dt.strftime('%Y%m%d')  # 修改为四位年份格式
        
        self.data.set_index('date', inplace=True)
        
        # 检查是否有 'code' 列
        if 'code' not in self.data.columns:
            if isinstance(self.data.index, pd.MultiIndex):
                self.data.reset_index(level='code', inplace=True)
            else:
                raise ValueError("数据中缺少'code'列")
    
    def process(self):
        """执行完整处理流程，按code分组计算指标"""
        result_dfs = []
        # 按code分组
        for code, group in self.data.groupby('code'):
            print(f"处理股票代码: {code}")
            # 确保group有date索引
            group = group.sort_index()
            # 计算买点信号
            calculator = TradeSignalCalculator(group)
            group_result = calculator.calculate_all_signals()
            result_dfs.append(group_result)
        
        # 合并所有结果
        self.data = pd.concat(result_dfs)
        
        # 只保留策略因子和必要的基础列
        strategy_factors = ['date', 'code', 'close', 'open', 'low', 'high', 'volume', 'ddx', 'POWERLINE','ENTRYBUY', 'BOTTOMBUY', 'BOTTOMUPBUY', 'POTENTIALBUY', 'STARK', 'DRAGONBUY', 'POWERDOWNSELL', 'CLEANSELL', 'STAGESELL', 'datekey']
        available_factors = [f for f in strategy_factors if f in self.data.columns]
        self.data = self.data[available_factors]
        
        return self.data
    
    def save_result(self, output_path):
        """保存结果，只包含策略因子"""
        self.data.to_parquet(output_path)

if __name__ == "__main__":
    # 使用示例
    # 替换为你的MySQL连接字符串
    db_conn_str = 'mysql+pymysql://root:root@192.168.44.137:3306/quant_factor'
    processor = TradeProcessor(db_connection_string=db_conn_str)
    result = processor.process()  
    print("处理完成，前5行结果:")
    print(result.head())
    processor.save_result('trade_indicators.parquet')
    #updated_df = process_buy_scores('trade_indicators.parquet')
    #print(updated_df[['code', 'datekey', 'buy_score']].head())